//
// network system calls.
//

#include "types.h"
#include "param.h"
#include "memlayout.h"
#include "riscv.h"
#include "spinlock.h"
#include "proc.h"
#include "defs.h"
#include "fs.h"
#include "sleeplock.h"
#include "file.h"
#include "net.h"

struct sock {
  struct sock *next; // the next socket in the list
  uint32 raddr;      // the remote IPv4 address
  uint16 lport;      // the local UDP port number
  uint16 rport;      // the remote UDP port number
  struct spinlock lock; // protects the rxq
  struct mbufq rxq;  // a queue of packets waiting to be received
  int transport_protocol; // transport's protocol, TCP or UDP
  struct tcpcb *tp;  // tcp连接信息
  uint32 seqnum;
  uint32 acknum;
};

static struct spinlock lock;
static struct sock *sockets;

void
sockinit(void)
{
  initlock(&lock, "socktbl");
}

int
sockalloc(struct file **f, uint32 raddr, uint16 lport, uint16 rport, int transport_protocol)
{
  struct sock *si, *pos;

  si = 0;
  *f = 0;
  if ((*f = filealloc()) == 0)
    goto bad;
  if ((si = (struct sock*)kalloc()) == 0)
    goto bad;

  // initialize objects
  si->raddr = raddr;
  si->lport = lport;
  si->rport = rport;
  si->transport_protocol = transport_protocol;
  si->tp = (struct tcpcb *)kalloc();
  if (si->tp == 0)
    goto bad;
  si->tp->t_state = TCPS_CLOSED;
  si->tp->t_receivednum = 0;
  si->seqnum = 0;
  si->acknum = 0;
  initlock(&si->lock, "sock");
  mbufq_init(&si->rxq);
  (*f)->type = FD_SOCK;
  (*f)->readable = 1;
  (*f)->writable = 1;
  (*f)->sock = si;

  // add to list of sockets
  acquire(&lock);
  pos = sockets;
  while (pos) {
    if (pos->raddr == raddr &&
        pos->lport == lport &&
	pos->rport == rport) {
      release(&lock);
      goto bad;
    }
    pos = pos->next;
  }
  si->next = sockets;
  sockets = si;
  release(&lock);
  return 0;

bad:
  if (si) {
    if (si->tp)
      kfree((char*)si->tp);
    kfree((char*)si);
  }
  if (*f)
    fileclose(*f);
  return -1;
}

void
sockclose(struct sock *si)
{
  struct sock **pos;
  struct mbuf *m;

  // remove from list of sockets
  acquire(&lock);
  pos = &sockets;
  while (*pos) {
    if (*pos == si){
      *pos = si->next;
      break;
    }
    pos = &(*pos)->next;
  }
  release(&lock);

  // 断开TCP连接
  si->tp->t_state = TCPS_FIN_WAIT_1;

  // free any pending mbufs
  while (!mbufq_empty(&si->rxq)) {
    m = mbufq_pophead(&si->rxq);
    mbuffree(m);
  }

  if (si->tp)
    kfree((char*)si->tp);
  kfree((char*)si);
}

int
sockread(struct sock *si, uint64 addr, int n)
{
  struct proc *pr = myproc();
  struct mbuf *m;
  int len;

  acquire(&si->lock);
  while (mbufq_empty(&si->rxq) && !pr->killed) {
    sleep(&si->rxq, &si->lock);
  }
  if (pr->killed) {
    release(&si->lock);
    return -1;
  }
  m = mbufq_pophead(&si->rxq);
  release(&si->lock);

  len = m->len;
  if (len > n)
    len = n;
  if (copyout(pr->pagetable, addr, m->head, len) == -1) {
    mbuffree(m);
    return -1;
  }
  mbuffree(m);
  return len;
}

int
sockwrite(struct sock *si, uint64 addr, int n)
{
  struct proc *pr = myproc();
  struct mbuf *m;

  if (si->transport_protocol == IPPROTO_TCP){
    uint32 seqnum, acknum;
    int send_total = 0;

    // 等待TCP连接建立
    acquire(&si->lock);
    while (si->tp->t_state<=TCPS_SYN_SENT && !pr->killed) {
      sleep(&si->rxq, &si->lock);
    }

    if (pr->killed) {
      release(&si->lock);
      return -1;
    }
    release(&si->lock);

    while(send_total<n) {
      acquire(&si->lock);
      while (si->tp->t_receivednum<send_total && !pr->killed) {
        sleep(&si->rxq, &si->lock);
      }

      if (pr->killed) {
        release(&si->lock);
        return -1;
      }

      seqnum = si->seqnum;
      acknum = si->acknum;
      release(&si->lock);

      m = mbufalloc(MBUF_DEFAULT_HEADROOM);
      if (!m)
        return -1;

      int send_size = TCP_MSS;
      if(send_total+send_size>n) {
        send_size = n-send_total;
      }

      if (copyin(pr->pagetable, mbufput(m, send_size), addr+send_total, send_size) == -1) {
        mbuffree(m);
        return -1;
      }

      if(si->tp->t_state==TCPS_SYN_RECEIVED){
        si->tp->t_state = TCPS_ESTABLISHED;
      }
      tcp_output(m, si->raddr, si->lport, si->rport, seqnum, acknum, TH_ACK, 1024);
      send_total += send_size;
    }

    // 等待TCP发送完成
    acquire(&si->lock);
    while (si->tp->t_receivednum<send_total && !pr->killed) {
      sleep(&si->rxq, &si->lock);
    }

    // 可以进行四次挥手中的两次挥手了，单方面断开自己的连接
    si->tp->t_state=TCPS_FIN_WAIT_1;
    m = mbufalloc(MBUF_DEFAULT_HEADROOM);
    // 除了带上TH_FIN，也需要带上TH_ACK以确认上次成功接收数据
    tcp_output(m, si->raddr, si->lport, si->rport, si->seqnum, si->acknum, TH_FIN|TH_ACK, 1024);

    if (pr->killed) {
      release(&si->lock);
      return -1;
    }
    release(&si->lock);

  } else { // IPPROTO_UDP 
    m = mbufalloc(MBUF_DEFAULT_HEADROOM);
    if (!m)
      return -1;

    if (copyin(pr->pagetable, mbufput(m, n), addr, n) == -1) {
      mbuffree(m);
      return -1;
    }
    net_tx_udp(m, si->raddr, si->lport, si->rport);
  }
  return n;
}

// called by protocol handler layer to deliver UDP packets
void
sockrecvudp(struct mbuf *m, uint32 raddr, uint16 lport, uint16 rport)
{
  //
  // Find the socket that handles this mbuf and deliver it, waking
  // any sleeping reader. Free the mbuf if there are no sockets
  // registered to handle it.
  //
  struct sock *si;

  acquire(&lock);
  si = sockets;
  while (si) {
    if (si->raddr == raddr && si->lport == lport && si->rport == rport)
      goto found;
    si = si->next;
  }
  release(&lock);
  mbuffree(m);
  return;

found:
  acquire(&si->lock);
  mbufq_pushtail(&si->rxq, m);
  wakeup(&si->rxq);
  release(&si->lock);
  release(&lock);
}

// called by protocol handler layer to deliver TCP packets
void
sockrecvtcp(struct tcp *tcphdr, struct mbuf *m, uint32 raddr, uint16 lport, uint16 rport, uint32 seqnum, uint32 acknum)
{
  struct sock *si;
  acquire(&lock);
  si = sockets;
  while (si) {
    if (si->raddr == raddr && si->lport == lport && 
          si->rport == rport)
      goto found;
    si = si->next;
  }
  release(&lock);
  mbuffree(m);
  return;

found:
  acquire(&si->lock);
  release(&lock);

  int queued = 0;
  /* handle SYN+ACK during handshake */
  if (si->tp->t_state == TCPS_SYN_SENT && (tcphdr->flags & TH_ACK) && (tcphdr->flags & TH_SYN)) {
    si->tp->t_state = TCPS_SYN_RECEIVED;
    si->acknum = seqnum+1;
  } else if (tcphdr->flags & TH_RST){
    mbuffree(m);
    release(&si->lock);
    return;
  }

  if(si->tp->t_state>=TCPS_ESTABLISHED && seqnum + m->len > si->acknum){
    si->acknum = seqnum + m->len;
    mbufq_pushtail(&si->rxq, m);
    queued = 1;
  }

  if (acknum > si->seqnum) {
    si->seqnum = acknum;
    si->tp->t_receivednum = acknum-1; // initial seq starts at 0, peer ACKs +1
  }

  if (!queued)
    mbuffree(m);

  wakeup(&si->rxq);
  release(&si->lock);
}

// tcp请求连接
int 
tcp_syn_connect(struct file *f)
{
  struct mbuf *m;

  m = mbufalloc(MBUF_DEFAULT_HEADROOM);
  if (!m)
    return -1;

  struct sock *si = f->sock;
  si->tp->t_state = TCPS_SYN_SENT;
  tcp_output(m, si->raddr, si->lport, si->rport, si->seqnum, si->acknum, TH_SYN, 1024);
  return 0;
}
